# 导入若干工具包
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim     # 导入优化器的包, optim中包含若干常用的优化算法, 比如SGD, Adam等
"""
1. 输入图片大小：（H,W）
2. 滤波器（卷积核）大小：（FH,FW），一般FH == FW
   填充(padding)为 P 
   步长为S 
3. 输出的图片大小：
   OH = (H + 2P -FH) /S +1
   OW = (W + 2P -FW) /S +1
   
池化的步长默认与核大小一样
实在把握不住尺寸，就拿一张图传入网络输出它的形状看看即可
"""

"""求解尺寸的 测试代码"""
a = torch.ones(3,32,32) # shape:(C,H,W)
# print(a)
b = torch.ones(1,1,32,32)   # 意思：一张1通道的32*32的图  # 必须注意格式，这是4维的数据
conv1 = nn.Conv2d(1, 6, 3)
conv2 = nn.Conv2d(6, 16, 3)
b1 = conv1(b)
b1 = F.max_pool2d(F.relu(b1), (2, 2))
b2 = conv2(b1)
b2 = F.max_pool2d(F.relu(b2), 2)
print(b2.size())            # 1 * 16 * 6 * 6
print("------------------------------------")
print("------------------------------------")

"""1. 定义一个简单的网络类"""
class Net(nn.Module):
    """使用的数据集输入的图片尺寸为32*32"""
    def __init__(self):
        super(Net, self).__init__()
        # 定义第一层卷积神经网络, 输入通道维度=1, 输出通道维度=6, 卷积核大小3*3
        self.conv1 = nn.Conv2d(1, 6, 3) # output.shape(30 * 30 * 6) 池化后 (15 * 15 * 6)
        # 定义第二层卷积神经网络, 输入通道维度=6, 输出通道维度=16, 卷积核大小3*3
        self.conv2 = nn.Conv2d(6, 16, 3) # output.shape(13 *13 * 16) 池化后 (6 * 6 * 16) 边缘不足2，默认去除，6.5 --> 6
        # 定义三层全连接网络
        self.fc1 = nn.Linear(16 * 6 * 6, 120)  # 全连接层第一次使用必须自行计算输入的像素的值
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # 在(2, 2)的池化窗口下执行最大池化操作
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)   # 步长默认与核大小一样
        x = x.view(-1, self.num_flat_features(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        # 计算size, 除了第0个维度上的batch_size
        size = x.size()[1:]
        num_features = 1
        for s in size:
            num_features *= s
        return num_features


net = Net()
print(net)
print("--------------------------")

# 模型中所有的可训练参数, 可以通过net.parameters()来获得.
# params = list(net.parameters())
# print(len(params))
# print(params[0].size())

# 测试网络
input1 = torch.randn(1, 1, 32, 32)
out = net(input1)
print(out)
print(out.size())

"""-*************************-"""
""" 2. 损失函数 """
"""
  损失函数的输入是一个输入的pair: (output, target), 然后计算出一个数值来评估output和target之间的差距大小.
  在torch.nn中有若干不同的损失函数可供使用, 比如nn.MSELoss就是通过计算均方差损失来评估输入和目标值之间的差距.
"""
output = net(input1)
target = torch.randn(10)

# 改变target的形状为二维张量, 为了和output匹配
target = target.view(1, -1)
criterion = nn.MSELoss()

loss = criterion(output, target)
print(loss)
"""     关于方向传播的链条: 如果我们跟踪loss反向传播的方向, 使用.grad_fn属性打印, 将可以看到一张完整的计算图如下:"""
"""
        input -> conv2d -> relu -> maxpool2d -> conv2d -> relu -> maxpool2d
              -> view -> linear -> relu -> linear -> relu -> linear
              -> MSELoss
              -> loss
"""
print(loss.grad_fn)  # MSELoss
print(loss.grad_fn.next_functions[0][0])     # Linear                  # 跟踪方法，往上去寻找上一个梯度对象
print(loss.grad_fn.next_functions[0][0].next_functions[0][0])  # ReLU  # 再往上去寻找更上一个梯度对象


"""-************************************************-"""
"""         3. 反向传播         """
"""
    在Pytorch中执行反向传播非常简便, 全部的操作就是loss.backward().
    在执行反向传播之前, 要先将梯度清零, 否则梯度会在不同的批次数据之间被累加.
"""
# Pytorch中执行梯度清零的代码
net.zero_grad()

print('conv1.bias.grad before backward')
print(net.conv1.bias.grad)   # 查看反向传播前第一个卷积层偏置值的梯度
# Pytorch中执行反向传播的代码
loss.backward()
print('conv1.bias.grad after backward')
print(net.conv1.bias.grad)   # 查看反向传播后第一个卷积层偏置值的梯度


"""-************************************************-"""
"""         4. 更新网络参数         """
"""
    更新参数最简单的算法就是SGD(随机梯度下降).
    具体的算法公式表达式为: weight = weight - learning_rate * gradient
"""
# 传统的python方法使用SGD更新
# learning_rate = 0.01
# for f in net.parameters():
#     f.data.sub_(f.grad.data * learning_rate)  # 手动梯度下降

# 深度学习框架中的参数更新方法
optimizer = optim.SGD(net.parameters(), lr=0.01) # 通过optim创建优化器对象

optimizer.zero_grad()                            # 将优化器执行梯度清零的操作       # 一个网络必须有的代码

output = net(input1)
loss = criterion(output, target)

loss.backward()                                  # 对损失值执行反向传播的操作       # 一个网络必须有的代码

optimizer.step()                                 # 参数的更新通过一行标准代码来执行  # 一个网络必须有的代码





